import pymesh
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 创建一个四面体圆柱网格
vertices = np.array([
    [0, 0, 0],
    [1, 0, 0],
    [0, 1, 0],
    [0, 0, 1],
    [0, 0, 2],
    [1, 0, 2],
    [0, 1, 2],
    [0, 0, 3]
])
faces = np.array([
    [0, 1, 2],
    [0, 1, 3],
    [0, 2, 3],
    [1, 2, 3],
    [4, 5, 6],
    [4, 5, 7],
    [4, 6, 7],
    [5, 6, 7],
    [0, 1, 4],
    [1, 4, 5],
    [1, 2, 5],
    [2, 5, 6],
    [2, 3, 6],
    [3, 6, 7],
    [3, 0, 7],
    [0, 4, 7]
])

mesh = pymesh.form_mesh(vertices, faces)

# 可视化四面体圆柱网格
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_trisurf(mesh.vertices[:, 0], mesh.vertices[:, 1], mesh.vertices[:, 2], triangles=mesh.faces)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()

# 输出网格的连接信息
for face in mesh.faces:
    print(f"Face: {face}")

# 两流体模型参数
# 假设你已经有了一些基本的物理参数
rho_l = 1000  # 液相密度 kg/m^3
rho_g = 1  # 气相密度 kg/m^3
mu_l = 0.001  # 液相粘度 Pa.s
mu_g = 1e-5  # 气相粘度 Pa.s
g = 9.81  # 重力加速度 m/s^2

# 初始化网格上的物理量
mesh.add_attribute("face_centroid")
mesh.add_attribute("face_area")
mesh.add_attribute("face_normal")

# 计算每个面的中心和面积
face_centroids = mesh.get_face_attribute("face_centroid")
face_areas = mesh.get_face_attribute("face_area")
face_normals = mesh.get_face_attribute("face_normal")

# 初始化两流体模型的变量
mesh.add_attribute("alpha_l")
mesh.add_attribute("alpha_g")
mesh.add_attribute("velocity_l")
mesh.add_attribute("velocity_g")
mesh.add_attribute("pressure")
mesh.add_attribute("enthalpy_l")
mesh.add_attribute("enthalpy_g")

# 假设初始条件
mesh.set_attribute("alpha_l", np.full(mesh.num_faces, 0.5))
mesh.set_attribute("alpha_g", np.full(mesh.num_faces, 0.5))
mesh.set_attribute("velocity_l", np.full(mesh.num_faces, 1.0))
mesh.set_attribute("velocity_g", np.full(mesh.num_faces, 0.1))
mesh.set_attribute("pressure", np.full(mesh.num_faces, 1e5))
mesh.set_attribute("enthalpy_l", np.full(mesh.num_faces, 1000))
mesh.set_attribute("enthalpy_g", np.full(mesh.num_faces, 2000))

# 定义控制容积方程
def mass_conservation(alpha_l, alpha_g, rho_l, rho_g):
    """
    质量守恒方程
    :param alpha_l: 液相体积分数
    :param alpha_g: 气相体积分数
    :param rho_l: 液相密度
    :param rho_g: 气相密度
    :return: 质量守恒结果
    """
    return alpha_l * rho_l + alpha_g * rho_g

def axial_momentum_conservation(alpha_l, alpha_g, rho_l, rho_g, velocity_l, velocity_g, pressure, g):
    """
    轴向动量守恒方程
    :param alpha_l: 液相体积分数
    :param alpha_g: 气相体积分数
    :param rho_l: 液相密度
    :param rho_g: 气相密度
    :param velocity_l: 液相速度
    :param velocity_g: 气相速度
    :param pressure: 压力
    :param g: 重力加速度
    :return: 轴向动量守恒结果
    """
    return alpha_l * rho_l * velocity_l + alpha_g * rho_g * velocity_g + pressure - rho_l * g

def lateral_momentum_conservation(alpha_l, alpha_g, rho_l, rho_g, velocity_l, velocity_g, pressure):
    """
    横向动量守恒方程
    :param alpha_l: 液相体积分数
    :param alpha_g: 气相体积分数
    :param rho_l: 液相密度
    :param rho_g: 气相密度
    :param velocity_l: 液相速度
    :param velocity_g: 气相速度
    :param pressure: 压力
    :return: 横向动量守恒结果
    """
    return alpha_l * rho_l * velocity_l**2 + alpha_g * rho_g * velocity_g**2 + pressure

def liquid_enthalpy_conservation(alpha_l, enthalpy_l, velocity_l):
    """
    液相焓守恒方程
    :param alpha_l: 液相体积分数
    :param enthalpy_l: 液相焓
    :param velocity_l: 液相速度
    :return: 液相焓守恒结果
    """
    return alpha_l * enthalpy_l * velocity_l

def gas_enthalpy_conservation(alpha_g, enthalpy_g, velocity_g):
    """
    气相焓守恒方程
    :param alpha_g: 气相体积分数
    :param enthalpy_g: 气相焓
    :param velocity_g: 气相速度
    :return: 气相焓守恒结果
    """
    return alpha_g * enthalpy_g * velocity_g

def mixture_enthalpy_conservation(alpha_l, alpha_g, enthalpy_l, enthalpy_g, velocity_l, velocity_g):
    """
    混合物焓守恒方程
    :param alpha_l: 液相体积分数
    :param alpha_g: 气相体积分数
    :param enthalpy_l: 液相焓
    :param enthalpy_g: 气相焓
    :param velocity_l: 液相速度
    :param velocity_g: 气相速度
    :return: 混合物焓守恒结果
    """
    return alpha_l * enthalpy_l * velocity_l + alpha_g * enthalpy_g * velocity_g


# 计算控制容积内的守恒量
for i in range(mesh.num_faces):
    alpha_l = mesh.get_attribute("alpha_l")[i]
    alpha_g = mesh.get_attribute("alpha_g")[i]
    velocity_l = mesh.get_attribute("velocity_l")[i]
    velocity_g = mesh.get_attribute("velocity_g")[i]
    pressure = mesh.get_attribute("pressure")[i]
    enthalpy_l = mesh.get_attribute("enthalpy_l")[i]
    enthalpy_g = mesh.get_attribute("enthalpy_g")[i]

    mass = mass_conservation(alpha_l, alpha_g, rho_l, rho_g)
    axial_momentum = axial_momentum_conservation(alpha_l, alpha_g, rho_l, rho_g, velocity_l, velocity_g, pressure, g)
    lateral_momentum = lateral_momentum_conservation(alpha_l, alpha_g, rho_l, rho_g, velocity_l, velocity_g, pressure)
    liquid_enthalpy = liquid_enthalpy_conservation(alpha_l, enthalpy_l, velocity_l)
    gas_enthalpy = gas_enthalpy_conservation(alpha_g, enthalpy_g, velocity_g)
    mixture_enthalpy = mixture_enthalpy_conservation(alpha_l, alpha_g, enthalpy_l, enthalpy_g, velocity_l, velocity_g)

    print(f"Face {i}:")
    print(f"Mass: {mass}")
    print(f"Axial Momentum: {axial_momentum}")
    print(f"Lateral Momentum: {lateral_momentum}")
    print(f"Liquid Enthalpy: {liquid_enthalpy}")
    print(f"Gas Enthalpy: {gas_enthalpy}")
    print(f"Mixture Enthalpy: {mixture_enthalpy}")
    print()